import torch as th
import numpy as np
import dgl

def random_aug(graph, attr, diag_attr, x, feat_drop_rate, edge_mask_rate):
    # 获取图中节点的数量
    n_node = graph.number_of_nodes()

    # 调用 mask_edge 函数，对图中的边进行掩码操作，得到未被掩码的边索引。
    edge_mask = mask_edge(graph, edge_mask_rate)
    # 对特征进行丢弃操作
    feat = drop_feature(x, feat_drop_rate)

    # 创建一个新的空图
    ng = dgl.graph([])
    # 向新图中添加节点
    ng.add_nodes(n_node)
    # 获取原图中的源节点和目标节点
    src = graph.edges()[0]
    dst = graph.edges()[1]
    # 根据掩码筛选源节点和目标节点
    nsrc = src[edge_mask]
    ndst = dst[edge_mask]
    # 向新图中添加边
    ng.add_edges(nsrc, ndst)
    # 向新图中添加自环边
    ng = ng.add_self_loop()

    # 更新边属性，将筛选后的边属性与对角线属性拼接在一起
    attr = th.cat([attr[edge_mask], diag_attr])
    return ng, attr, feat

def drop_feature(x, drop_prob):
    # 生成一个与特征维度相同的随机掩码，根据 drop_prob 进行随机丢弃
    drop_mask = th.empty(
        (x.size(1),),
        dtype=th.float32,
        device=x.device).uniform_(0, 1) < drop_prob
    # 克隆输入特征
    x = x.clone()
    # 将掩码位置的特征值置为 0
    x[:, drop_mask] = 0

    return x

def mask_edge(graph, mask_prob):
    # 获取图中边的数量
    E = graph.number_of_edges()
    # 生成一个大小为 E 的掩码概率张量，值为 mask_prob
    mask_rates = th.FloatTensor(np.ones(E) * mask_prob)
    # 根据掩码概率生成伯努利分布的掩码
    masks = th.bernoulli(1 - mask_rates)
    # 获取未被掩码的位置索引
    mask_idx = masks.nonzero().squeeze(1)
    return mask_idx
